from typing import Dict, Any, List
import json
import requests
import time
from casevo import LLM_INTERFACE
from chromadb import Documents, EmbeddingFunction, Embeddings

class BaichuanEmbedding(EmbeddingFunction):
    """百川嵌入函数"""
    
    def __init__(self, llm, tar_len):
        self.SEND_LEN = tar_len
        self.llm = llm
    
    def __call__(self, input: Documents) -> Embeddings:
        """计算文档嵌入"""
        res_list = []
        cur_list = []
        for item in input:
            cur_list.append(item)
            if len(cur_list) >= self.SEND_LEN:
                res = self.llm.send_embedding(cur_list)
                res_list.extend(res)
                cur_list = []
        if len(cur_list) > 0:
            res = self.llm.send_embedding(cur_list)
            res_list.extend(res)
        return res_list

class BaichuanLLM(LLM_INTERFACE):
    """百川LLM实现"""
    
    def __init__(self, config: Dict[str, Any]):
        self.api_key = config.get('api_key')
        if not self.api_key:
            raise ValueError("API key is required")
            
        self.headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        self.embedding_len = config.get('embedding_len', 10)
        self.embedding_function = BaichuanEmbedding(self, self.embedding_len)
        
    def generate(self, prompt: str) -> str:
        """生成响应"""
        return self.send_message(prompt, json_flag=True)
        
    def send_message(self, prompt: str, json_flag: bool = False) -> str:
        """发送消息"""
        print('发送提示词:', prompt)
        
        data = {
            'model': 'Baichuan3-Turbo',
            'messages': [{'role': 'user', 'content': prompt}],
            'temperature': 0.3,
            'top_p': 0.85,
            'max_tokens': 2048,
            'with_search_enhance': False
        }
        
        try:
            response = requests.post(
                'https://api.baichuan-ai.com/v1/chat/completions',
                headers=self.headers,
                json=data
            )
            
            # 处理速率限制
            while response.status_code == 429:
                print('API速率限制，等待10秒...')
                time.sleep(5)
                response = requests.post(
                    'https://api.baichuan-ai.com/v1/chat/completions',
                    headers=self.headers,
                    json=data
                )
                
            if response.status_code == 200:
                result = json.loads(response.text)
                
                # 处理内容过滤
                if result['choices'][0]['finish_reason'] == 'content_filter':
                    print('内容被过滤，重试...')
                    time.sleep(5)
                    response = requests.post(
                        'https://api.baichuan-ai.com/v1/chat/completions',
                        headers=self.headers,
                        json=data
                    )
                    result = json.loads(response.text)
                    
                content = result['choices'][0]['message']['content'].strip()
                print('收到响应:', content)
                return content
            else:
                print(f'API调用失败: {response.status_code}')
                return "{}"
                
        except Exception as e:
            print(f'API调用异常: {str(e)}')
            return "{}"
            
    def send_embedding(self, text_list: List[str]) -> List[List[float]]:
        """发送嵌入请求"""
        data = {
            'model': 'Baichuan-Text-Embedding',
            'input': text_list
        }
        
        try:
            response = requests.post(
                'http://api.baichuan-ai.com/v1/embeddings',
                headers=self.headers,
                json=data
            )
            
            # 处理速率限制
            while response.status_code == 429:
                print('API速率限制，等待10秒...')
                time.sleep(5)
                response = requests.post(
                    'http://api.baichuan-ai.com/v1/embeddings',
                    headers=self.headers,
                    json=data
                )
                
            if response.status_code == 200:
                result = json.loads(response.text)
                return [item['embedding'] for item in result['data']]
            else:
                print(f'嵌入API调用失败: {response.status_code}')
                return []
                
        except Exception as e:
            print(f'嵌入API调用异常: {str(e)}')
            return []
            
    def get_lang_embedding(self) -> EmbeddingFunction:
        """获取语言嵌入函数"""
        return self.embedding_function 